from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
# from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
import os

# from dashscope import DashScopeEmbedding
# from llama_index.embeddings.dashscope import DashScopeEmbedding

# from langchain_community.embeddings import OpenAIEmbeddings
# 假设阿里百炼的嵌入类名为 DashScopeEmbeddings

# from llama_index.embeddings.dashscope import (
#     DashScopeEmbedding,
#     DashScopeTextEmbeddingModels,
#     DashScopeTextEmbeddingType,
# )
# from llama_index.embeddings.dashscope import (
#     DashScopeEmbedding,
#     DashScopeTextEmbeddingModels,
#     DashScopeTextEmbeddingType,
# )
from langchain_community.embeddings import DashScopeEmbeddings


os.environ["OPENAI_API_KEY"] = "sk-f5324346ba744ef89eda093af8f307c7"
os.environ["OPENAI_API_BASE"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
os.environ["DASHSCOPE_API_KEY"] = "sk-f5324346ba744ef89eda093af8f307c7"

loader = TextLoader("introduction.txt", encoding="UTF-8")
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)

# 使用 DashScopeEmbeddings
# embedding_function = DashScopeEmbeddings(api_key=os.environ["OPENAI_API_KEY"])
# embedder = OpenAIEmbeddings(model="multimodal-embedding-v1",
#                             openai_api_base='https://dashscope.aliyuncs.com/compatible-mode/v1',
#                             openai_api_key='sk-f5324346ba744ef89eda093af8f307c7')
# embedder = DashScopeEmbeddings(
#     model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V2,
#     text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
# )
# embedder = DashScopeEmbeddings(
#     model_name="multimodal-embedding-v1",
#     dashscope_api_key="sk-f5324346ba744ef89eda093af8f307c7",
# )

embedder = DashScopeEmbeddings(
    model="text-embedding-v3",
    dashscope_api_key=os.environ["DASHSCOPE_API_KEY"],
)

vectorstore = Chroma(
    collection_name="ai_learning",
    embedding_function=embedder,
    persist_directory="vectordb"
)
vectorstore.add_documents(splits)
